#lang scheme

(require
 (lib "1.ss" "srfi")
 (lib "13.ss" "srfi")
 (lib "14.ss" "srfi")
 (lib "26.ss" "srfi") ;;cut
 (lib "process.ss") ;;open/save
 (lib "defmacro.ss")
 "deBruijn.scm"
 "linus.scm"
 "state.scm"
 "utils.scm"
 "transition.scm")

(provide fst%)
(define fst%
  (class
   object%
   (define R 25)
   (define states '())
   (define transitions '())
   (define inputs '())
   (define outputs '())
   (define input-history '())
   (define output-history '())
   (define cur-state '())
   (define start-state '())
   (define continuation-point '*)
   (define amb-fail (lambda x #f))
   (define paths (make-hash))
   (define sequence-type 'DeBruijn)

   (define-macro amb
     (lambda alts...
       `(let ((+prev-amb-fail amb-fail))
          (call/cc
           (lambda (+sk)
             ,@(map (lambda (alt)
                      `(call/cc
                        (lambda (+fk)
                          (set! amb-fail
                                (lambda ()
                                  (set! amb-fail +prev-amb-fail)
                                  (+fk)))
                          (+sk ,alt))))
                    alts...)
             (+prev-amb-fail))))))


   (define/public (get-neighbors source fst)
     (map (cut get-field tail <>)
          (filter (lambda (t)
                    (eq? source
                         (get-field head t)))
                  (send fst get-transitions))))


   (define/public (fst-equal? fst2)
     (for*/or
      ((s1 states)
       (s2 (send fst2 get-states)))
      (let loop ((fringe1 `(,s1))
                 (fringe2 `(,s2))
                 (closed1 '())
                 (closed2 '()))
        (if (or (null? fringe1) (null? fringe2))
            (and (= (length closed1) (length closed2))
                 ;;(printd "closed1 = ~s~n" (map (cut send <> get-name) closed1))
                 ;;(printd "closed2 = ~s~n" (map (cut send <> get-name) closed2))
                 (for/and
                  ((s1 closed1)
                   (s2 closed2)
                   #:when #t
                   (i inputs))
                  (for*/or
                   ((t1 transitions)
                    (t2 (send fst2 get-transitions)))
                   (and (equal? (get-field head t1) s1)
                        (equal? (get-field head t2) s2)
                        (equal? (get-field input t1) i)
                        (equal? (get-field input t2) i)
                        (equal? (get-field output t1)
                                (get-field output t2))
                        (eq? (list-index (cut equal? <> (get-field tail t1)) closed1)
                             (list-index (cut equal? <> (get-field tail t2)) closed2))))))
            (loop (append (cdr fringe1)
                          (lset-difference
                           equal?
                           (send this get-neighbors (car fringe1) this)
                           fringe1
                           closed1))
                  (append (cdr fringe2)
                          (lset-difference
                           equal?
                           (send this get-neighbors (car fringe2) fst2)
                           fringe2
                           closed2))
                  (cons (car fringe1) closed1)
                  (cons (car fringe2) closed2))))))


   (define (an-element-of items)
     (if (null? items) (amb)
         (amb (car items) (an-element-of (cdr items)))))


   ;;returns the transition that leads to dest or #f is there is not path there
   (define/public (dijkstra source dest)
;     (printd "source = ~s~n" (send source get-name))
;     (printd "dest = ~s~n" (send dest get-name))
     (let ((src-row (hash-ref paths source)))
       (unless
        (hash-ref src-row dest #f)
        (let loop ((fringe `(,source))
                   (closed '()))
          (when
           (pair? fringe)
           (let ([neighbors (lset-difference
                             equal?
                             (send this get-neighbors (car fringe) this)
                             fringe
                             closed)])
             (for ((x neighbors))
                  (hash-set! src-row x (hash-ref src-row (car fringe) x)))
             (loop (append (cdr fringe) neighbors)
                   (cons (car fringe) closed)))))
        (hash-set! src-row source source))
       (find (lambda (t)
               (and (equal? (get-field tail t) (hash-ref src-row dest #f))
                    (equal? (get-field head t) source)))
             transitions)))

   (init-field parent)

   (define/public (open file)
     (parameterize
      ([current-input-port (open-input-file file)])
      (read-line)
      (read-line)
      (match-let
       ([(list _ neato-width neato-height)
         (regexp-match #px"(\\d+)\\.?\\d*,(\\d+)\\.?\\d*\"" (read-line))])
       (for ([line (in-lines)])
            (match line
                   [(regexp #px"(\\S+) -> (\\S+) \\S+ (\\S+) (\\S+).*"
                            (list _ head tail input output))
                    'foo
                    (add-transition (find (lambda (s) (equal? (send s get-name)
                                                              head))
                                          states)
                                    (find (lambda (s) (equal? (send s get-name)
                                                              tail))
                                          states)
                                    input output)]
                   [(regexp #px"(\\S+).*\"(\\d+)\\.?\\d*,(\\d+)\\.?\\d*.*"
                            (list _ head x y))
                    (add-state (* (string->number x)
                                  (/ (send parent get-width)
                                     (string->number neato-width)))
                               (* (string->number y)
                                  (/ (send parent get-height)
                                     (string->number neato-height)))
                               head)]
                   [else #f])))
      (close-input-port (current-input-port))))


   (define/public (save file)
     (parameterize
      ([current-output-port
        (open-output-file file #:exists 'truncate)])
      (printf "digraph g {~n")
      (for ([s states])
           (printf "~a;~n" (send s get-name)))
      (for ([t transitions])
           (let ((in (get-field input t))
                 (out (get-field output t)))
             (printf "~a -> ~a [label=\" ~a ~a \"]~n"
                     (send (get-field head t) get-name)
                     (send (get-field tail t) get-name)
                     (if in in "?")
                     (if out out "?"))))
      (printf "}~n")
      (close-output-port (current-output-port)))
     (system (format "neato ~s -o ~s 2>/dev/null" file file)))


   (define/public (consume-input c)
     (for/first ([t transitions]
                 #:when (equal? (get-field head t) cur-state)
                 #:when (equal? (get-field input t)
                                (cond [(char? c) (list->string (list c))]
                                      [(symbol? c ) (symbol->string c)]
                                      [else c])))
                (cons-last! c input-history)
                (cons-last! (get-field output t) output-history)
                (send cur-state consume-input c)
                (set-cur-state! (get-field tail t))
                (send parent refresh)
                (get-field output t)))


   (define/public (next-step)
     (let ((right-fst (get-field fst (cadr (send (send parent get-parent) get-children))))
           (dest-state (send this least-visited)))
       (printd "state with smallest history is ~s~n" (send dest-state get-name))
       (let* ((symbol (if (equal? cur-state dest-state)
                            (next-symbol (send cur-state get-history) inputs)
                            (get-field input (send this dijkstra cur-state dest-state))))
              (out-left (send this consume-input symbol))
              (out-right (send right-fst consume-input symbol)))
         (set! input-history (send right-fst get-input-history))
         (set! output-history (send right-fst get-output-history))
         (printd "Consumed ~s on both left and right~n" symbol)
         (printd "left-output ~s~n" out-left)
         (printd "right-output ~s~n" out-right)
         (printd "input-history ~s~n" input-history)
         (printd "output-history ~s~n" output-history)
         (for ((s states))
              (printd "history of state ~s: ~s~n"
                      (send s get-name)
                      (send s get-history)))
         (unless (equal? out-left out-right)
                 (printd "outputs differ: generating new fst~n")
                 (call-with-continuation-prompt
                  continuation-point
                  (default-continuation-prompt-tag))
                 (for ((s states))
                      (printd "history of state ~s: ~s~n"
                              (send s get-name)
                              (send s get-history)))
                 (printf "input history = ~s~n" input-history)
                 (printf "fst-equal? ~s~s~n"
                         (send this fst-equal?
                               (get-field
                                fst
                                (cadr (send (send parent get-parent) get-children))))
                         (length input-history))))))


   (define/public (next-fst)
     (printd "increasing size of fst~n")
     (cons-last! (new state% (name (get-new-state-name states)) (x 0) (y 0))
                 states)
     (let loop ((i 0)
                (new-transitions '())
                (new-cur-state (car states)))
       (call/cc
        (lambda (cc)
          (set! continuation-point cc)))
       (when (= i (length input-history))
             (printd "updating fst to new guess~n")
             (set! transitions new-transitions)

             ;; find missing transitions
             (for* ((i inputs)
                    (s states)
                    #:when (not
                            (find
                             (lambda (t) (and (equal? (send (get-field head t) get-name) (send s get-name))
                                              (equal? (get-field input t) i)))
                             transitions)))
                   (printd "adding transition ~s->~s:~s/~s~n"
                           (send s get-name)
                           (send (car states) get-name)
                           i
                           (car outputs))
                   (add-transition s (car states) i (car outputs)))

             ;; save and reopen to have dotty display them nicely
             (save TEMP_FILE_NAME)
             (set! states '())
             (set! transitions '())
             (open TEMP_FILE_NAME)

             ;; clear the paths matrix because all the transitions have chanegd
             (set! paths (make-hash))
             (for ((s states)) (hash-set! paths s (make-hash)))

             ;; make sure the resulting machine is fully connected
             (unless (send this fully-connected?)
                     (printd "not fully-connected~n")
                     (amb))

             ;; run the input through the machine to verify it and fill back in the histories
             (set! input-history '())
             (set! output-history '())
             ;; doing so duplicates the history on the right, so reset that as well.
             (let ((right-fst (get-field fst (cadr (send (send parent get-parent) get-children)))))
               (for ((s (send right-fst get-input-history)))
                    (send this consume-input s))
               (send right-fst set-input-history! input-history)
               (send right-fst set-output-history! output-history))

             ;; return the guess machine
             (abort-current-continuation
              (default-continuation-prompt-tag)
              (lambda () 0)))
       (printd "~s->?:~s/~s~n"
               (send new-cur-state get-name)
               (list-ref input-history i)
               (list-ref output-history i))
       (let  ((t (find (lambda (t)
                         (and (equal? (get-field input t)
                                      (list-ref input-history i))
                              (equal? (send (get-field head t) get-name)
                                      (send new-cur-state get-name))))
                       new-transitions)))
         (cond [(not t)
                (let ((tail (an-element-of states)))
                  (unless tail
                          (printd "No more backtracking options~n")
                          (next-fst))
;                  (printd "Transition table is:~n")
;                  (for ((t new-transitions))
;                       (printd "~s->~s:~s/~s~n"
;                               (send (get-field head t) get-name)
;                               (send (get-field tail t) get-name)
;                               (get-field input t)
;                               (get-field output t)))
;                  (printd "~s->~s:~s/~s~n"
;                          (send new-cur-state get-name)
;                          (send tail get-name)
;                          (list-ref input-history i)
;                          (list-ref output-history i))
;                  (printd "-----~n")
                  (loop (+ i 1)
                        (append
                         new-transitions
                         (list
                          (new transition%
                               (head new-cur-state) (tail tail)
                               (input (list-ref input-history i))
                               (output (list-ref output-history i)))))
                        tail))]
               [(equal? (list-ref output-history i) (get-field output t))
                (loop (+ i 1) new-transitions (get-field tail t))]
               [else
                (printd "Backtracking...~n")
                (amb)]))))


   (define (get-new-state-name states)
     (find (lambda (n)
             (not (find (lambda (s) (equal? (send s get-name) n))
                        states)))
           (map list->string (zip (map integer->char (iota 26 65))))))


   (define/public (next-symbol history inputs)
     (printf "   sequence type = ~s~n" sequence-type)
     (case sequence-type
       ((Random) (list-ref inputs (random (length inputs))))
       ((Linus) (next-linus-symbol history inputs))
       ((DeBruijn) (next-deBruijn-symbol history inputs))))


   (define/public (add-state x y (name (get-input "name")))
     (let ((s (new state% (name name) (x x) (y y))))
       (when (null? states)
             (set-start-state! s)
             (set-cur-state! s))
       (cons-last! s states)
       s))


   (define/public (add-transition head tail
                                  (input (get-input "input symbol"))
                                  (output (get-input "output symbol")))
     (unless (member input inputs)
             (add-input-symbol input))
     (unless (member output outputs)
             (add-output-symbol output))
     (cons-last! (new transition%
                 (head head) (tail tail)
                 (input input) (output output)) transitions)
     (let ((t (find (lambda (t) (and (eq? (get-field head t) tail)
                                     (eq? (get-field tail t) head)))
                    transitions)))
       (when t (send t set-arc) (send (last transitions) set-arc))))


   (define/public (least-visited)
     (let-values
      (((state hist-len)
        (for/fold ((best-state cur-state)
                   (best-hist-len (length (send cur-state get-history))))
                  ((s states))
                  (let ((h (send s get-history)))
                    (if (length< h best-hist-len)
                        (values s (length h))
                        (values best-state best-hist-len))))))
      state))


   (define (random-element l)
     (list-ref l (random (length l))))

   (define/public (create-random-fst nodes input-alphabet output-alphabet)
     (set! states '())
     (set! transitions '())
     (set! input-history '())
     (set! output-history '())
     (set! inputs input-alphabet)
     (set! outputs output-alphabet)

     ;; create the states
     (for ((n nodes)) (add-state 0 0 n))

     ;; create random transitions
     (for* ((s states)
            (i inputs))
           (add-transition s (random-element states)
                           i (random-element outputs)))

     ;; clear the paths matrix
     (set! paths (make-hash))
     (for ((s states)) (hash-set! paths s (make-hash)))

     ;;and then check to see if it's fully connected
     (printd "checking to see if fully connected~n")
     (unless (send this fully-connected?)
      (create-random-fst nodes input-alphabet output-alphabet)))


   (define/public (fully-connected?)
     (for*/and ((s states)
                (d states)
                #:when (not (equal? d s)))
               (send this dijkstra s d)))


   (define/public (add-input-symbol symbol) (cons-last! symbol inputs))
   (define/public (add-output-symbol symbol) (cons-last! symbol outputs))
   (define/public (get-cur-state) cur-state)
   (define/public (get-inputs) inputs)
   (define/public (get-input-history) input-history)
   (define/public (get-outputs) outputs)
   (define/public (get-output-history) output-history)
   (define/public (get-paths) paths)
   (define/public (get-start-state) start-state)
   (define/public (get-states) states)
   (define/public (get-transitions) transitions)
   (define/public (set-input-history! s) (set! input-history s))
   (define/public (set-output-history! s) (set! output-history s))
   (define/public (set-start-state! s) (set! start-state s))
   (define/public (set-cur-state! s) (set! cur-state s))
   (define/public (set-sequence-type! t) (set! sequence-type t))
   (super-new)))
